{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "accbccab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from streamz import Stream\n",
    "import hvplot.streamz\n",
    "from streamz.river import RiverTrain\n",
    "from river import cluster\n",
    "import holoviews as hv\n",
    "from panel.pane.holoviews import HoloViews\n",
    "hv.extension('bokeh')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a2ef27a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = cluster.KMeans(n_clusters=3, sigma=0.1, mu=0.5)\n",
    "centres = [[random.random(), random.random()] for _ in range(3)]\n",
    "\n",
    "def gen(move_chance=0.05):\n",
    "    centre = int(random.random() * 3)  # 3x faster than random.randint(0, 2)\n",
    "    if random.random() < move_chance:\n",
    "        centres[centre][0] += random.random() / 5 - 0.1\n",
    "        centres[centre][1] += random.random() / 5 - 0.1\n",
    "    value = {'x': random.random() / 20 + centres[centre][0],\n",
    "             'y': random.random() / 20 + centres[centre][1]}\n",
    "    return value\n",
    "\n",
    "\n",
    "def get_clusters(model):\n",
    "    # return [{\"x\": xcen, \"y\": ycen}, ...] for each centre\n",
    "    data = [{'x': v['x'], 'y': v['y']} for k, v in model.centers.items()]\n",
    "    return pd.DataFrame(data, index=range(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6451048",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = Stream.from_periodic(gen, 0.03)\n",
    "km = RiverTrain(model, pass_model=True)\n",
    "s.map(lambda x: (x,)).connect(km)  # learn takes a tuple of (x,[ y[, w]])\n",
    "ex = pd.DataFrame({'x': [0.5], 'y': [0.5]})\n",
    "ooo = s.map(lambda x: pd.DataFrame([x])).to_dataframe(example=ex)\n",
    "out = km.map(get_clusters)\n",
    "\n",
    "# start things\n",
    "s.emit(gen())  # set initial model\n",
    "for i, (x, y) in enumerate(centres):\n",
    "    model.centers[i]['x'] = x\n",
    "    model.centers[i]['y'] = y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b4de451",
   "metadata": {},
   "outputs": [],
   "source": [
    "pout = out.to_dataframe(example=ex)\n",
    "pl = (ooo.hvplot.scatter('x', 'y', color=\"blue\", backlog=50) *\n",
    "      pout.hvplot.scatter('x', 'y', color=\"red\", backlog=3))\n",
    "pl.opts(xlim=(-0.2, 1.2), ylim=(-0.2, 1.2), height=600, width=600)\n",
    "pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c24d2363",
   "metadata": {},
   "outputs": [],
   "source": [
    "s.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18cfd94e",
   "metadata": {},
   "outputs": [],
   "source": [
    "s.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4537495c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
